// Tencent is pleased to support the open source community by making RapidJSON available.
// 
// Copyright (C) 2015 THL A29 Limited, a Tencent company, and Milo Yip. All rights reserved.
//
// Licensed under the MIT License (the "License"); you may not use this file except
// in compliance with the License. You may obtain a copy of the License at
//
// http://opensource.org/licenses/MIT
//
// Unless required by applicable law or agreed to in writing, software distributed 
// under the License is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR 
// CONDITIONS OF ANY KIND, either express or implied. See the License for the 
// specific language governing permissions and limitations under the License.

#ifndef RAPIDJSON_INTERNAL_REGEX_H_
#define RAPIDJSON_INTERNAL_REGEX_H_

#include "../allocators.h"
#include "../stream.h"
#include "stack.h"

#ifdef __clang__
RAPIDJSON_DIAG_PUSH
RAPIDJSON_DIAG_OFF(padded)
RAPIDJSON_DIAG_OFF(switch-
                           enum)
RAPIDJSON_DIAG_OFF(implicit - fallthrough)
#endif

#ifdef __GNUC__
RAPIDJSON_DIAG_PUSH
RAPIDJSON_DIAG_OFF(effc++)
#endif

#ifdef _MSC_VER
RAPIDJSON_DIAG_PUSH
RAPIDJSON_DIAG_OFF(4512) // assignment operator could not be generated
#endif

#ifndef RAPIDJSON_REGEX_VERBOSE
#define RAPIDJSON_REGEX_VERBOSE 0
#endif

RAPIDJSON_NAMESPACE_BEGIN
    namespace internal {

///////////////////////////////////////////////////////////////////////////////
// GenericRegex

        static const SizeType kRegexInvalidState = ~SizeType(
                0);  //!< Represents an invalid index in GenericRegex::State::out, out1
        static const SizeType kRegexInvalidRange = ~SizeType(0);

//! Regular expression engine with subset of ECMAscript grammar.
/*!
    Supported regular expression syntax:
    - \c ab     Concatenation
    - \c a|b    Alternation
    - \c a?     Zero or one
    - \c a*     Zero or more
    - \c a+     One or more
    - \c a{3}   Exactly 3 times
    - \c a{3,}  At least 3 times
    - \c a{3,5} 3 to 5 times
    - \c (ab)   Grouping
    - \c ^a     At the beginning
    - \c a$     At the end
    - \c .      Any character
    - \c [abc]  Character classes
    - \c [a-c]  Character class range
    - \c [a-z0-9_] Character class combination
    - \c [^abc] Negated character classes
    - \c [^a-c] Negated character class range
    - \c [\b]   Backspace (U+0008)
    - \c \\| \\\\ ...  Escape characters
    - \c \\f Form feed (U+000C)
    - \c \\n Line feed (U+000A)
    - \c \\r Carriage return (U+000D)
    - \c \\t Tab (U+0009)
    - \c \\v Vertical tab (U+000B)

    \note This is a Thompson NFA engine, implemented with reference to 
        Cox, Russ. "Regular Expression Matching Can Be Simple And Fast (but is slow in Java, Perl, PHP, Python, Ruby,...).", 
        https://swtch.com/~rsc/regexp/regexp1.html 
*/
        template<typename Encoding, typename Allocator = CrtAllocator>
        class GenericRegex {
        public:
            typedef typename Encoding::Ch Ch;

            GenericRegex(const Ch *source, Allocator *allocator = 0) :
                    states_(allocator, 256), ranges_(allocator, 256), root_(kRegexInvalidState), stateCount_(),
                    rangeCount_(),
                    stateSet_(), state0_(allocator, 0), state1_(allocator, 0), anchorBegin_(), anchorEnd_() {
                GenericStringStream<Encoding> ss(source);
                DecodedStream <GenericStringStream<Encoding>> ds(ss);
                Parse(ds);
            }

            ~GenericRegex() {
                Allocator::Free(stateSet_);
            }

            bool IsValid() const {
                return root_ != kRegexInvalidState;
            }

            template<typename InputStream>
            bool Match(InputStream &is) const {
                return SearchWithAnchoring(is, true, true);
            }

            bool Match(const Ch *s) const {
                GenericStringStream<Encoding> is(s);
                return Match(is);
            }

            template<typename InputStream>
            bool Search(InputStream &is) const {
                return SearchWithAnchoring(is, anchorBegin_, anchorEnd_);
            }

            bool Search(const Ch *s) const {
                GenericStringStream<Encoding> is(s);
                return Search(is);
            }

        private:
            enum Operator {
                kZeroOrOne,
                kZeroOrMore,
                kOneOrMore,
                kConcatenation,
                kAlternation,
                kLeftParenthesis
            };

            static const unsigned kAnyCharacterClass = 0xFFFFFFFF;   //!< For '.'
            static const unsigned kRangeCharacterClass = 0xFFFFFFFE;
            static const unsigned kRangeNegationFlag = 0x80000000;

            struct Range {
                unsigned start; //
                unsigned end;
                SizeType next;
            };

            struct State {
                SizeType out;     //!< Equals to kInvalid for matching state
                SizeType out1;    //!< Equals to non-kInvalid for split
                SizeType rangeStart;
                unsigned codepoint;
            };

            struct Frag {
                Frag(SizeType s, SizeType o, SizeType m) : start(s), out(o), minIndex(m) {}

                SizeType start;
                SizeType out; //!< link-list of all output states
                SizeType minIndex;
            };

            template<typename SourceStream>
            class DecodedStream {
            public:
                DecodedStream(SourceStream &ss) : ss_(ss), codepoint_() { Decode(); }

                unsigned Peek() { return codepoint_; }

                unsigned Take() {
                    unsigned c = codepoint_;
                    if (c) // No further decoding when '\0'
                        Decode();
                    return c;
                }

            private:
                void Decode() {
                    if (!Encoding::Decode(ss_, &codepoint_))
                        codepoint_ = 0;
                }

                SourceStream &ss_;
                unsigned codepoint_;
            };

            State &GetState(SizeType index) {
                RAPIDJSON_ASSERT(index < stateCount_);
                return states_.template Bottom<State>()[index];
            }

            const State &GetState(SizeType index) const {
                RAPIDJSON_ASSERT(index < stateCount_);
                return states_.template Bottom<State>()[index];
            }

            Range &GetRange(SizeType index) {
                RAPIDJSON_ASSERT(index < rangeCount_);
                return ranges_.template Bottom<Range>()[index];
            }

            const Range &GetRange(SizeType index) const {
                RAPIDJSON_ASSERT(index < rangeCount_);
                return ranges_.template Bottom<Range>()[index];
            }

            template<typename InputStream>
            void Parse(DecodedStream<InputStream> &ds) {
                Allocator allocator;
                Stack<Allocator> operandStack(&allocator, 256);     // Frag
                Stack<Allocator> operatorStack(&allocator, 256);    // Operator
                Stack<Allocator> atomCountStack(&allocator, 256);   // unsigned (Atom per parenthesis)

                *atomCountStack.template Push<unsigned>() = 0;

                unsigned codepoint;
                while (ds.Peek() != 0) {
                    switch (codepoint = ds.Take()) {
                        case '^':
                            anchorBegin_ = true;
                            break;

                        case '$':
                            anchorEnd_ = true;
                            break;

                        case '|':
                            while (!operatorStack.Empty() && *operatorStack.template Top<Operator>() < kAlternation)
                                if (!Eval(operandStack, *operatorStack.template Pop<Operator>(1)))
                                    return;
                            *operatorStack.template Push<Operator>() = kAlternation;
                            *atomCountStack.template Top<unsigned>() = 0;
                            break;

                        case '(':
                            *operatorStack.template Push<Operator>() = kLeftParenthesis;
                            *atomCountStack.template Push<unsigned>() = 0;
                            break;

                        case ')':
                            while (!operatorStack.Empty() &&
                                   *operatorStack.template Top<Operator>() != kLeftParenthesis)
                                if (!Eval(operandStack, *operatorStack.template Pop<Operator>(1)))
                                    return;
                            if (operatorStack.Empty())
                                return;
                            operatorStack.template Pop<Operator>(1);
                            atomCountStack.template Pop<unsigned>(1);
                            ImplicitConcatenation(atomCountStack, operatorStack);
                            break;

                        case '?':
                            if (!Eval(operandStack, kZeroOrOne))
                                return;
                            break;

                        case '*':
                            if (!Eval(operandStack, kZeroOrMore))
                                return;
                            break;

                        case '+':
                            if (!Eval(operandStack, kOneOrMore))
                                return;
                            break;

                        case '{': {
                            unsigned n, m;
                            if (!ParseUnsigned(ds, &n))
                                return;

                            if (ds.Peek() == ',') {
                                ds.Take();
                                if (ds.Peek() == '}')
                                    m = kInfinityQuantifier;
                                else if (!ParseUnsigned(ds, &m) || m < n)
                                    return;
                            } else
                                m = n;

                            if (!EvalQuantifier(operandStack, n, m) || ds.Peek() != '}')
                                return;
                            ds.Take();
                        }
                            break;

                        case '.':
                            PushOperand(operandStack, kAnyCharacterClass);
                            ImplicitConcatenation(atomCountStack, operatorStack);
                            break;

                        case '[': {
                            SizeType range;
                            if (!ParseRange(ds, &range))
                                return;
                            SizeType s = NewState(kRegexInvalidState, kRegexInvalidState, kRangeCharacterClass);
                            GetState(s).rangeStart = range;
                            *operandStack.template Push<Frag>() = Frag(s, s, s);
                        }
                            ImplicitConcatenation(atomCountStack, operatorStack);
                            break;

                        case '\\': // Escape character
                            if (!CharacterEscape(ds, &codepoint))
                                return; // Unsupported escape character
                            // fall through to default

                        default: // Pattern character
                            PushOperand(operandStack, codepoint);
                            ImplicitConcatenation(atomCountStack, operatorStack);
                    }
                }

                while (!operatorStack.Empty())
                    if (!Eval(operandStack, *operatorStack.template Pop<Operator>(1)))
                        return;

                // Link the operand to matching state.
                if (operandStack.GetSize() == sizeof(Frag)) {
                    Frag *e = operandStack.template Pop<Frag>(1);
                    Patch(e->out, NewState(kRegexInvalidState, kRegexInvalidState, 0));
                    root_ = e->start;

#if RAPIDJSON_REGEX_VERBOSE
                    printf("root: %d\n", root_);
                    for (SizeType i = 0; i < stateCount_ ; i++) {
                        State& s = GetState(i);
                        printf("[%2d] out: %2d out1: %2d c: '%c'\n", i, s.out, s.out1, (char)s.codepoint);
                    }
                    printf("\n");
#endif
                }

                // Preallocate buffer for SearchWithAnchoring()
                RAPIDJSON_ASSERT(stateSet_ == 0);
                if (stateCount_ > 0) {
                    stateSet_ = static_cast<unsigned *>(states_.GetAllocator().Malloc(GetStateSetSize()));
                    state0_.template Reserve<SizeType>(stateCount_);
                    state1_.template Reserve<SizeType>(stateCount_);
                }
            }

            SizeType NewState(SizeType out, SizeType out1, unsigned codepoint) {
                State *s = states_.template Push<State>();
                s->out = out;
                s->out1 = out1;
                s->codepoint = codepoint;
                s->rangeStart = kRegexInvalidRange;
                return stateCount_++;
            }

            void PushOperand(Stack<Allocator> &operandStack, unsigned codepoint) {
                SizeType s = NewState(kRegexInvalidState, kRegexInvalidState, codepoint);
                *operandStack.template Push<Frag>() = Frag(s, s, s);
            }

            void ImplicitConcatenation(Stack<Allocator> &atomCountStack, Stack<Allocator> &operatorStack) {
                if (*atomCountStack.template Top<unsigned>())
                    *operatorStack.template Push<Operator>() = kConcatenation;
                (*atomCountStack.template Top<unsigned>())++;
            }

            SizeType Append(SizeType l1, SizeType l2) {
                SizeType old = l1;
                while (GetState(l1).out != kRegexInvalidState)
                    l1 = GetState(l1).out;
                GetState(l1).out = l2;
                return old;
            }

            void Patch(SizeType l, SizeType s) {
                for (SizeType next; l != kRegexInvalidState; l = next) {
                    next = GetState(l).out;
                    GetState(l).out = s;
                }
            }

            bool Eval(Stack<Allocator> &operandStack, Operator op) {
                switch (op) {
                    case kConcatenation:
                        RAPIDJSON_ASSERT(operandStack.GetSize() >= sizeof(Frag) * 2);
                        {
                            Frag e2 = *operandStack.template Pop<Frag>(1);
                            Frag e1 = *operandStack.template Pop<Frag>(1);
                            Patch(e1.out, e2.start);
                            *operandStack.template Push<Frag>() = Frag(e1.start, e2.out, Min(e1.minIndex, e2.minIndex));
                        }
                        return true;

                    case kAlternation:
                        if (operandStack.GetSize() >= sizeof(Frag) * 2) {
                            Frag e2 = *operandStack.template Pop<Frag>(1);
                            Frag e1 = *operandStack.template Pop<Frag>(1);
                            SizeType s = NewState(e1.start, e2.start, 0);
                            *operandStack.template Push<Frag>() = Frag(s, Append(e1.out, e2.out),
                                                                       Min(e1.minIndex, e2.minIndex));
                            return true;
                        }
                        return false;

                    case kZeroOrOne:
                        if (operandStack.GetSize() >= sizeof(Frag)) {
                            Frag e = *operandStack.template Pop<Frag>(1);
                            SizeType s = NewState(kRegexInvalidState, e.start, 0);
                            *operandStack.template Push<Frag>() = Frag(s, Append(e.out, s), e.minIndex);
                            return true;
                        }
                        return false;

                    case kZeroOrMore:
                        if (operandStack.GetSize() >= sizeof(Frag)) {
                            Frag e = *operandStack.template Pop<Frag>(1);
                            SizeType s = NewState(kRegexInvalidState, e.start, 0);
                            Patch(e.out, s);
                            *operandStack.template Push<Frag>() = Frag(s, s, e.minIndex);
                            return true;
                        }
                        return false;

                    default:
                        RAPIDJSON_ASSERT(op == kOneOrMore);
                        if (operandStack.GetSize() >= sizeof(Frag)) {
                            Frag e = *operandStack.template Pop<Frag>(1);
                            SizeType s = NewState(kRegexInvalidState, e.start, 0);
                            Patch(e.out, s);
                            *operandStack.template Push<Frag>() = Frag(e.start, s, e.minIndex);
                            return true;
                        }
                        return false;
                }
            }

            bool EvalQuantifier(Stack<Allocator> &operandStack, unsigned n, unsigned m) {
                RAPIDJSON_ASSERT(n <= m);
                RAPIDJSON_ASSERT(operandStack.GetSize() >= sizeof(Frag));

                if (n == 0) {
                    if (m == 0)                             // a{0} not support
                        return false;
                    else if (m == kInfinityQuantifier)
                        Eval(operandStack, kZeroOrMore);    // a{0,} -> a*
                    else {
                        Eval(operandStack, kZeroOrOne);         // a{0,5} -> a?
                        for (unsigned i = 0; i < m - 1; i++)
                            CloneTopOperand(operandStack);      // a{0,5} -> a? a? a? a? a?
                        for (unsigned i = 0; i < m - 1; i++)
                            Eval(operandStack, kConcatenation); // a{0,5} -> a?a?a?a?a?
                    }
                    return true;
                }

                for (unsigned i = 0; i < n - 1; i++)        // a{3} -> a a a
                    CloneTopOperand(operandStack);

                if (m == kInfinityQuantifier)
                    Eval(operandStack, kOneOrMore);         // a{3,} -> a a a+
                else if (m > n) {
                    CloneTopOperand(operandStack);          // a{3,5} -> a a a a
                    Eval(operandStack, kZeroOrOne);         // a{3,5} -> a a a a?
                    for (unsigned i = n; i < m - 1; i++)
                        CloneTopOperand(operandStack);      // a{3,5} -> a a a a? a?
                    for (unsigned i = n; i < m; i++)
                        Eval(operandStack, kConcatenation); // a{3,5} -> a a aa?a?
                }

                for (unsigned i = 0; i < n - 1; i++)
                    Eval(operandStack, kConcatenation);     // a{3} -> aaa, a{3,} -> aaa+, a{3.5} -> aaaa?a?

                return true;
            }

            static SizeType Min(SizeType a, SizeType b) { return a < b ? a : b; }

            void CloneTopOperand(Stack<Allocator> &operandStack) {
                const Frag src = *operandStack.template Top<Frag>(); // Copy constructor to prevent invalidation
                SizeType count = stateCount_ -
                                 src.minIndex; // Assumes top operand contains states in [src->minIndex, stateCount_)
                State *s = states_.template Push<State>(count);
                memcpy(s, &GetState(src.minIndex), count * sizeof(State));
                for (SizeType j = 0; j < count; j++) {
                    if (s[j].out != kRegexInvalidState)
                        s[j].out += count;
                    if (s[j].out1 != kRegexInvalidState)
                        s[j].out1 += count;
                }
                *operandStack.template Push<Frag>() = Frag(src.start + count, src.out + count, src.minIndex + count);
                stateCount_ += count;
            }

            template<typename InputStream>
            bool ParseUnsigned(DecodedStream<InputStream> &ds, unsigned *u) {
                unsigned r = 0;
                if (ds.Peek() < '0' || ds.Peek() > '9')
                    return false;
                while (ds.Peek() >= '0' && ds.Peek() <= '9') {
                    if (r >= 429496729 && ds.Peek() > '5') // 2^32 - 1 = 4294967295
                        return false; // overflow
                    r = r * 10 + (ds.Take() - '0');
                }
                *u = r;
                return true;
            }

            template<typename InputStream>
            bool ParseRange(DecodedStream<InputStream> &ds, SizeType *range) {
                bool isBegin = true;
                bool negate = false;
                int step = 0;
                SizeType start = kRegexInvalidRange;
                SizeType current = kRegexInvalidRange;
                unsigned codepoint;
                while ((codepoint = ds.Take()) != 0) {
                    if (isBegin) {
                        isBegin = false;
                        if (codepoint == '^') {
                            negate = true;
                            continue;
                        }
                    }

                    switch (codepoint) {
                        case ']':
                            if (start == kRegexInvalidRange)
                                return false;   // Error: nothing inside []
                            if (step == 2) { // Add trailing '-'
                                SizeType r = NewRange('-');
                                RAPIDJSON_ASSERT(current != kRegexInvalidRange);
                                GetRange(current).next = r;
                            }
                            if (negate)
                                GetRange(start).start |= kRangeNegationFlag;
                            *range = start;
                            return true;

                        case '\\':
                            if (ds.Peek() == 'b') {
                                ds.Take();
                                codepoint = 0x0008; // Escape backspace character
                            } else if (!CharacterEscape(ds, &codepoint))
                                return false;
                            // fall through to default

                        default:
                            switch (step) {
                                case 1:
                                    if (codepoint == '-') {
                                        step++;
                                        break;
                                    }
                                    // fall through to step 0 for other characters

                                case 0: {
                                    SizeType r = NewRange(codepoint);
                                    if (current != kRegexInvalidRange)
                                        GetRange(current).next = r;
                                    if (start == kRegexInvalidRange)
                                        start = r;
                                    current = r;
                                }
                                    step = 1;
                                    break;

                                default:
                                    RAPIDJSON_ASSERT(step == 2);
                                    GetRange(current).end = codepoint;
                                    step = 0;
                            }
                    }
                }
                return false;
            }

            SizeType NewRange(unsigned codepoint) {
                Range *r = ranges_.template Push<Range>();
                r->start = r->end = codepoint;
                r->next = kRegexInvalidRange;
                return rangeCount_++;
            }

            template<typename InputStream>
            bool CharacterEscape(DecodedStream<InputStream> &ds, unsigned *escapedCodepoint) {
                unsigned codepoint;
                switch (codepoint = ds.Take()) {
                    case '^':
                    case '$':
                    case '|':
                    case '(':
                    case ')':
                    case '?':
                    case '*':
                    case '+':
                    case '.':
                    case '[':
                    case ']':
                    case '{':
                    case '}':
                    case '\\':
                        *escapedCodepoint = codepoint;
                        return true;
                    case 'f':
                        *escapedCodepoint = 0x000C;
                        return true;
                    case 'n':
                        *escapedCodepoint = 0x000A;
                        return true;
                    case 'r':
                        *escapedCodepoint = 0x000D;
                        return true;
                    case 't':
                        *escapedCodepoint = 0x0009;
                        return true;
                    case 'v':
                        *escapedCodepoint = 0x000B;
                        return true;
                    default:
                        return false; // Unsupported escape character
                }
            }

            template<typename InputStream>
            bool SearchWithAnchoring(InputStream &is, bool anchorBegin, bool anchorEnd) const {
                RAPIDJSON_ASSERT(IsValid());
                DecodedStream<InputStream> ds(is);

                state0_.Clear();
                Stack<Allocator> *current = &state0_, *next = &state1_;
                const size_t stateSetSize = GetStateSetSize();
                std::memset(stateSet_, 0, stateSetSize);

                bool matched = AddState(*current, root_);
                unsigned codepoint;
                while (!current->Empty() && (codepoint = ds.Take()) != 0) {
                    std::memset(stateSet_, 0, stateSetSize);
                    next->Clear();
                    matched = false;
                    for (const SizeType *s = current->template Bottom<SizeType>();
                         s != current->template End<SizeType>(); ++s) {
                        const State &sr = GetState(*s);
                        if (sr.codepoint == codepoint ||
                            sr.codepoint == kAnyCharacterClass ||
                            (sr.codepoint == kRangeCharacterClass && MatchRange(sr.rangeStart, codepoint))) {
                            matched = AddState(*next, sr.out) || matched;
                            if (!anchorEnd && matched)
                                return true;
                        }
                        if (!anchorBegin)
                            AddState(*next, root_);
                    }
                    internal::Swap(current, next);
                }

                return matched;
            }

            size_t GetStateSetSize() const {
                return (stateCount_ + 31) / 32 * 4;
            }

            // Return whether the added states is a match state
            bool AddState(Stack<Allocator> &l, SizeType index) const {
                RAPIDJSON_ASSERT(index != kRegexInvalidState);

                const State &s = GetState(index);
                if (s.out1 != kRegexInvalidState) { // Split
                    bool matched = AddState(l, s.out);
                    return AddState(l, s.out1) || matched;
                } else if (!(stateSet_[index >> 5] & (1 << (index & 31)))) {
                    stateSet_[index >> 5] |= (1 << (index & 31));
                    *l.template PushUnsafe<SizeType>() = index;
                }
                return s.out ==
                       kRegexInvalidState; // by using PushUnsafe() above, we can ensure s is not validated due to reallocation.
            }

            bool MatchRange(SizeType rangeIndex, unsigned codepoint) const {
                bool yes = (GetRange(rangeIndex).start & kRangeNegationFlag) == 0;
                while (rangeIndex != kRegexInvalidRange) {
                    const Range &r = GetRange(rangeIndex);
                    if (codepoint >= (r.start & ~kRangeNegationFlag) && codepoint <= r.end)
                        return yes;
                    rangeIndex = r.next;
                }
                return !yes;
            }

            Stack<Allocator> states_;
            Stack<Allocator> ranges_;
            SizeType root_;
            SizeType stateCount_;
            SizeType rangeCount_;

            static const unsigned kInfinityQuantifier = ~0u;

            // For SearchWithAnchoring()
            uint32_t *stateSet_;        // allocated by states_.GetAllocator()
            mutable Stack<Allocator> state0_;
            mutable Stack<Allocator> state1_;
            bool anchorBegin_;
            bool anchorEnd_;
        };

        typedef GenericRegex<UTF8<> > Regex;

    } // namespace internal
RAPIDJSON_NAMESPACE_END

#ifdef __clang__
RAPIDJSON_DIAG_POP
#endif

#ifdef _MSC_VER
RAPIDJSON_DIAG_POP
#endif

#endif // RAPIDJSON_INTERNAL_REGEX_H_
